import argparse
import os
import numpy as np
import tensorflow as tf
from tensorflow.keras import Sequential, layers
parser = argparse.ArgumentParser(description="training setting.")
parser.add_argument("--data_path", type=str, default='./',
                    help="data_path.")
parser.add_argument("--batch_size", type=int, default=100,
                    help="batch_size.")
parser.add_argument("--output_path", type=str, default='./',
                    help="output_path.")
args_input = parser.parse_args()

if not os.path.exists(args_input.output_path):
    os.mkdir(args_input.output_path)
tensorboard = tf.keras.callbacks.TensorBoard(log_dir=args_input.output_path)

with np.load(os.path.join(args_input.data_path, "MNIST-data-0")) as f:
    trainImage, trainLabel = f['x_train'], f['y_train']
trainImage = tf.reshape(trainImage, (60000, 28, 28, 1))
network = Sequential([
    layers.Conv2D(filters=6, kernel_size=(5, 5), activation="relu", input_shape=(28, 28, 1), padding="same"),
    layers.MaxPool2D(pool_size=(2, 2), strides=2),
    layers.Conv2D(filters=16, kernel_size=(5, 5), activation="relu", padding="same"),
    layers.MaxPool2D(pool_size=2, strides=2),
    layers.Conv2D(filters=32, kernel_size=(5, 5), activation="relu", padding="same"),
    layers.Flatten(),
    layers.Dense(200, activation="relu"),
    layers.Dense(10, activation="softmax")
])
network.summary()

network.compile(optimizer='adam', loss="sparse_categorical_crossentropy", metrics=["accuracy"])
network.fit(trainImage, trainLabel, epochs=1, validation_split=0.1, callbacks=[tensorboard])

network.save(os.path.join(args_input.output_path, 'lenet_mnist.h5'))
print('lenet_mnist model saved')
del network
